﻿using StableDiffusionTools.Oobabooga;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace LLMEvaluation.Framework
{
    public class GenerationRun
    {
        public int Runs { get; set; } = 1;
        public int Seed = -1;
        public List<Message> PromptMessages { get; set; } = new List<Message>();
        public List<ModelEntry> Models { get; set; } = new List<ModelEntry>();
        public RunEntry Run { get; set; }
        public bool HasModels { get => Models.Count > 0; }
        public bool AllModelsIdle
        {
            get
            {
                bool result = true;

                //TODO: track idle models by counter, should be more performant.
                foreach (var model in Models)
                {
                    if(model.Inference != InferenceState.Idle)
                    {
                        result = false;
                        break;
                    }
                }

                return result;
            }
        }

        public GenerationRun(RunEntry run) : this(run, run.PromptMessages, run.ModelEntries) { }
        public GenerationRun(RunEntry run, List<ModelEntry> models) : this(run, run.PromptMessages, models) { }

        public GenerationRun(RunEntry runModel, List<Message> prompt, List<ModelEntry> modelEntries)
        {
            Run = runModel;

            Runs = runModel.Runs;

            for (int i = 0; i < prompt.Count; i++)
                PromptMessages.Add(new Message(prompt[i].Role, Message.FixToApi(prompt[i].Content)));

            Models = new List<ModelEntry>(modelEntries.Where((m) => !m.Name.Equals("None")));

            Seed = Run.Seed;
        }

        public static List<string> GetUniqueModelNames(List<GenerationRun> runs)
        {
            List<string> models = new List<string>();

            for (int i = 0; i < runs.Count; i++)
            {
                var run = runs[i];
                for (int j = 0; j < run.Models.Count; j++)
                {
                    if (!models.Contains(run.Models[j].Name))
                        models.Add(run.Models[j].Name);
                }
            }
            return models;
        }

        public static List<string> GetUniqueModelNames(List<RunEntry> runs)
        {
            List<string> models = new List<string>();

            for (int i = 0; i < runs.Count; i++)
            {
                var run = runs[i];
                for (int j = 0; j < run.ModelEntries.Count; j++)
                {
                    if (!models.Contains(run.ModelEntries[j].Name))
                        models.Add(run.ModelEntries[j].Name);
                }
            }
            return models;
        }

        public void Schedule() => SetInferenceState(InferenceState.Scheduled);
        public void Unschedule() => SetInferenceState(InferenceState.Idle);

        public void SetInferenceState(InferenceState state)
        {
            Run.Inference = state;
            for (int i = 0; i < Models.Count; i++)
            {
                Models[i].Inference = state;
            }
        }
    }
}
